import shutil
import random
import tkinter as tk
from  tkinter import filedialog, ttk, messagebox, font

from PIL import Image, ImageTk

import os
# import cv2

from pathlib import Path


COLOR_LIST = [
    'old lace',
    'linen', 'blanched almond', 'bisque', 'peach puff',
    'navajo white', 'lemon chiffon', 'mint cream', 'azure', 'alice blue', 'lavender',
    'lavender blush', 'misty rose', 
    'light gray', 'cornflower blue',
    'light slate blue', 'royal blue',
    'dodger blue', 'deep sky blue', 'sky blue', 'steel blue', 'light steel blue',
    'light blue', 'dark turquoise', 'medium turquoise', 'turquoise',
    'cyan','cadet blue', 'medium aquamarine', 'aquamarine',
    'dark sea green', 'medium sea green', 'light sea green', 'pale green', 'spring green',
    'lawn green', 'medium spring green', 'green yellow', 'lime green', 'yellow green',
    'forest green', 'olive drab', 'dark khaki', 'khaki', 'pale goldenrod', 'light goldenrod yellow',
    'yellow', 'gold', 'goldenrod', 'dark goldenrod', 'rosy brown',
    'indian red', 'sandy brown',
    'dark salmon', 'salmon', 'light salmon', 'orange', 
    'coral', 'light coral', 'tomato', 'orange red', 'hot pink', 'pink',
    'pale violet red', 'medium violet red',
    'medium orchid', 'medium purple',
    'thistle'
]

class App(tk.Frame):
    def __init__(self, master):
        super().__init__(master)
        self.master.minsize(1280, 720)
        self.master.maxsize(self.winfo_screenwidth(), self.winfo_screenheight())
        self.place(relx=0, rely=0, relwidth=1, relheight=1)

        self.num_img = 0
        self.current_ind = 0
        self.current_cls = None
        self.quick_mode = False
        self.fresh_rest_list = False
        self.to_be_deleted_inds = []

        self.key_binds = [
            '1', '2', '3', '4', '5',
            'Q', 'W', 'E', 'R', 'T',
            'A', 'S', 'D', 'F', 'G',
            'Z', 'X', 'C', 'V', 'B'
        ]

        self.contents_open_path = tk.StringVar()
        self.contents_save_path = tk.StringVar()

        lf_path = tk.LabelFrame(self, text='设置')
        lf_path.place(relx=0.01, rely=0.01, relwidth=0.98, relheight=0.1)
        self.lf_path = lf_path

        self.butt_open_path = tk.Button(lf_path)
        self.butt_open_path['text'] = "打开文件夹"
        self.butt_open_path['command'] = self._open_dir
        self.butt_open_path.place(relx=0.0, rely=0.0, relwidth=0.1, relheight=0.48)

        label_op = tk.Label(lf_path, bg='white')
        label_op["textvariable"] = self.contents_open_path
        label_op.place(relx=0.1, relwidth=0.75, rely=0.0, relheight=0.48)

        label_sp = tk.Label(lf_path, bg='white')
        label_sp["textvariable"] = self.contents_save_path
        label_sp.place(relx=0.1, relwidth=0.75, rely=0.49, relheight=0.48)

        self.butt_save_csv = tk.Button(lf_path)
        self.butt_save_csv['text'] = "导出csv文件到"
        self.butt_save_csv['command'] = self._save_csv_file
        self.butt_save_csv.place(relx=0.0, relwidth=0.1, rely=0.49, relheight=0.48)

        self.butt_open_cls_file = tk.Button(lf_path)
        self.butt_open_cls_file['text'] = "一键按类别文件夹导出"
        self.butt_open_cls_file['command'] = self._save_dir
        self.butt_open_cls_file.place(relx=0.85, relwidth=0.15, rely=0.49, relheight=0.48)


        self.cmb_fmt = ttk.Combobox(lf_path, takefocus=False)
        self.cmb_fmt['value'] = ('.jpg', '.png')
        self.cmb_fmt.place(relx=0.85, relwidth=0.08, rely=0.0, relheight=0.48)
        # self.cmb_fmt.state(['readonly'])
        self.cmb_fmt.current(0)

        self.butt_pause = tk.Button(lf_path, text="设置类型", command=self._set_fmt)
        # self.butt_pause['state'] = 'disabled'
        self.butt_pause.place(relx=0.93, relwidth=0.07, rely=0.0, relheight=0.48)

        # 类别设置
        self.classnames = []
        if Path('predefined_classes.txt').exists():
            with open('predefined_classes.txt') as f:
                for line in f:
                    self.classnames.append(line.strip())
        if len(self.classnames) > 20:
            messagebox.showwarning(
                '警告！', 
                '仅支持最多20类别设置！\n目前有{}个类别，选择前20个'.format(len(self.classnames))
            )
            self.classnames = self.classnames[:20]
        
        self._set_class_buttions()

        self.lf_image = tk.LabelFrame(self, text="图像以及标注")
        self.lf_image.place(relx=0.22, rely=0.29, relwidth=0.77, relheight=0.7)

        self.lb_image = tk.Label(self.lf_image, image=None)  # light yellow
        self.lb_image.place(relx=0.0, rely=0.0, relwidth=1, relheight=1)

        self.master.bind("<KeyPress>", self._keypress)
        # self.bind("<Button-1>", self.click)

        lf_list_all = tk.LabelFrame(self, text='所有列表')
        lf_list_all.place(relx=0.01, relwidth=0.2, rely=0.29, relheight=0.4)
        
        scroll_all = tk.Scrollbar(lf_list_all)
        scroll_all.pack(side=tk.RIGHT, fill=tk.Y)
        scroll_all.place(relx=0.9, relwidth=0.1, rely=0, relheight=1)

        self.list_all = tk.Listbox(lf_list_all, yscrollcommand=scroll_all.set)
        self.list_all.place(relx=0, relwidth=0.9, rely=0, relheight=1)
        self.list_all.bind("<<ListboxSelect>>", self._click_on_list)
        scroll_all.config(command=self.list_all.yview)


        #  text='未标注列表'
        self.butt_fresh_rest = tk.Button(self)
        self.butt_fresh_rest['text'] = "点击刷新未标注列表"
        self.butt_fresh_rest['command'] = self._fresh_rest_list
        self.butt_fresh_rest.place(relx=0.01, relwidth=0.09, rely=0.7, relheight=0.03)

        lf_list_rest = tk.LabelFrame(self)
        lf_list_rest.place(relx=0.01, relwidth=0.2, rely=0.73, relheight=0.26)
        
        scroll_rest = tk.Scrollbar(lf_list_rest)
        scroll_rest.pack(side=tk.RIGHT, fill=tk.Y)
        scroll_rest.place(relx=0.9, relwidth=0.1, rely=0, relheight=1)

        self.contents_rest_list = tk.StringVar()
        self.list_rest = tk.Listbox(
            lf_list_rest, 
            yscrollcommand=scroll_rest.set,
            listvariable=self.contents_rest_list
        )
        self.list_rest.place(relx=0, relwidth=0.9, rely=0, relheight=1)
        self.list_rest.bind("<<ListboxSelect>>", self._click_on_list_rest)
        scroll_rest.config(command=self.list_rest.yview)
        
        # tk.Radiobutton(master)
    

    def _fresh_rest_list(self):
        # rest_paths = eval(self.contents_rest_list.get())
        rest_paths = sorted(self.rel_paths_set)
        self.contents_rest_list.set(r' '.join(rest_paths))


    def _set_class_buttions(self):

        interval = len(COLOR_LIST) // len(self.classnames)
        # print(interval)
        self.selected_colors = COLOR_LIST[::interval][:len(self.classnames)]

        self.lf_classname = tk.LabelFrame(self, text='类别设置 @普通模式 (按“空格”健 进入 快速模式)')
        self.lf_classname.place(relx=0.01, rely=0.12, relwidth=0.98, relheight=0.16)

        self.var_cls = tk.IntVar()
        # self.var_cls.set(-1)

        num_cls = len(self.classnames)
        max_num_col = 5
        if num_cls > max_num_col:
            inter_w = 1 / max_num_col
            inter_h = 1 / (1 + (num_cls - 1) // max_num_col)
        else:
            inter_w = 1 / num_cls
            inter_h = 1

        self.rb_dict = {}
        for ind, name in enumerate(self.classnames):
            self.rb_dict[name] = tk.Radiobutton(
                self.lf_classname, 
                text='(' + self.key_binds[ind] + ') ' + str(name),
                font=font.Font(family='Consolas', size=9),
                bg=self.selected_colors[ind],
                anchor=tk.W,
                variable=self.var_cls, value=ind,
                command=self._click_to_classify
                # indicatoron=0
            )
            self.rb_dict[name].select()
            
            self.rb_dict[name].place(
                relx=(ind % max_num_col) * inter_w, relwidth=inter_w-0.001, 
                rely=(ind // max_num_col) * inter_h, relheight=inter_h-0.05
            )
    
    # def _open_class_file(self):
    #     self.classnames = []
    #     with open(filedialog.askopenfilename(initialdir='./')) as f:
    #         for line in f:
    #             self.classnames.append(line.strip())
    #     self._set_class_buttions()
    #     self.butt_open_cls_file['state'] = 'disabled'

    def _click_to_classify(self):
        # print(self.classnames[self.var_cls.get()])
        self.all_labels[self.current_ind] = self.var_cls.get()
        self.current_cls = self.var_cls.get()
        self._show_image()

    def _set_fmt(self):
        self.cmb_fmt['state'] = 'normal'
    
    def _click_on_list(self, event):
        self.current_ind = event.widget.curselection()[0]
        self._show_image()
    
    def _click_on_list_rest(self, event):
        ind = event.widget.curselection()[0]
        self.current_ind = self.rel_paths.index(eval(self.contents_rest_list.get())[ind])
        self._show_image()

    def _open_dir(self):
        self.open_path = Path(filedialog.askdirectory())
        self.contents_open_path.set(str(self.open_path))
        
        self.cmb_fmt['state'] = 'disabled'
        # print(self.cmb_fmt.get())
        self.all_img_paths = sorted(self.open_path.glob('**/*{}'.format(self.cmb_fmt.get())))
        self.num_img = len(self.all_img_paths)

        if self.num_img == 0:
            messagebox.showinfo('提示', '未找到格式为{}的图片!\n请修改格式重新打开文件夹！'.format(self.cmb_fmt.get()))
            self._set_fmt()
        else:
            messagebox.showinfo(
                '提示', 
                '已找到{}张图片！\n请使用类别对应案件进行打标签分类！\n请使用 “←” “→” 进行翻页！\nENJOY!'.format(self.num_img)
            )

        self.fresh_rest_list = False
        self.list_all.delete(0, tk.END)
        self.to_be_deleted_inds = []
        self.rel_paths = [str(p.relative_to(self.open_path).as_posix()) for p in self.all_img_paths]
        self.rel_paths_set = set(self.rel_paths)
        for p in self.rel_paths:
            self.list_all.insert(tk.END, p)
        
        self.contents_rest_list.set(' '.join(self.rel_paths))
            # self.list_rest.insert(tk.END, p)
        # self.rest_paths = self.list_rest.get(0, tk.END)
        self.all_labels = [None, ] * self.num_img
        
        self._show_image()
    

    def _save_dir(self):
        ret = True
        if len(self.to_be_deleted_inds) < self.num_img:
            ret = messagebox.askyesno('警告', '还未完成标注，确定复制到分类文件夹吗？\n未完成的标注将不保存', icon=messagebox.WARNING)
        if ret:
            save_path = Path(filedialog.askdirectory(
                initialdir=str(self.open_path.parent)
            )) / (self.open_path.stem + '_picked')

            self.contents_save_path.set(str(save_path))
            if save_path.exists():
                ret = messagebox.askyesno('警告', '已存在文件夹，是否删除？', icon=messagebox.WARNING)
                if ret:
                    shutil.rmtree(save_path)
            else:
                messagebox.showinfo("提示", "将自动创建 {} 文件夹".format(self.open_path.stem + '_picked'))
            save_path.mkdir(exist_ok=True, parents=True)

            for name in self.classnames:
                (save_path / name).mkdir(exist_ok=True)

            prog = ttk.Progressbar(self.lf_path, maximum=len(self.to_be_deleted_inds))
            prog.place(relx=0.1, relwidth=0.75, rely=0.49, relheight=0.48)
            prog['value'] = 0

            for ind, p in enumerate(self.all_img_paths):
                if self.all_labels[ind] is not None:
                    prog['value'] += 1
                    dst_p = str(save_path / self.classnames[self.all_labels[ind]] / p.name)
                    shutil.copy(p, dst_p)
                    self.update()
                    # f.write('{},{}\n'.format(str(p), self.classnames[self.all_labels[ind]]))

            messagebox.showinfo("提示", "图片导出完成！")
            prog.destroy()

    def _save_csv_file(self):
        ret = True
        if len(self.to_be_deleted_inds) < self.num_img:
            ret = messagebox.askyesno('警告', '还未完成标注，确定导出csv文件吗？\n未完成的标注将不保存', icon=messagebox.WARNING)
        if ret:
            save_path = Path(filedialog.asksaveasfilename(
                initialdir=str(self.open_path.parent),
                initialfile=self.open_path.stem + '.csv'
            ))
            self.contents_save_path.set(str(save_path))
            save_path.parent.mkdir(exist_ok=True, parents=True)
            with open(save_path, 'w') as f:
                for ind, p in enumerate(self.all_img_paths):
                    if self.all_labels[ind] is not None:
                        f.write('{},{}\n'.format(str(p), self.classnames[self.all_labels[ind]]))
            
            messagebox.showinfo("提示", "csv文件导出完成！")


    def _keypress(self, event):
        # print(self.contents_open_path.get())
        if self.num_img:
            # print(event.keysym)
            
            if event.keysym.upper() in self.key_binds:
                self.current_cls = self.key_binds.index(event.keysym.upper())
                self.all_labels[self.current_ind] = self.current_cls
                self.var_cls.set(self.current_cls)

            if event.keysym in ['Escape']:
                if self.quick_mode:
                    messagebox.showinfo(
                        '提示', 
                        '退出快速标注模式！不再继承已选择类别！'
                    )
                    self.lf_classname['text'] = '类别设置 @普通模式 (按“空格”健 进入 快速模式)'
                self.current_cls = None
                self.quick_mode = False
            
            if event.keysym in ['space']:
                if not self.quick_mode:
                    messagebox.showinfo(
                        '提示', 
                        '启用快速标注模式！\n本模式下，鼠标点击或者按键选择的类别将具有继承性，\n翻页继承当前选择的类别！'
                    )
                    self.lf_classname['text'] = '类别设置 @快速模式 (按“ESC”健 进入 普通模式)'
                self.quick_mode = True
                

            if event.keysym in ['Left']:
                # print('previous')
                self.current_ind -= 1
                if self.current_ind < 0:
                    self.current_ind = 0
            elif event.keysym in ['Right']:
                # print('next')
                self.current_ind += 1
                if self.current_ind > self.num_img - 1:
                    self.current_ind = self.num_img - 1
        
            if self.quick_mode:
                self.all_labels[self.current_ind] = self.current_cls

            # self.list_all.activate(self.current_ind)

            self._show_image()
    
    def _show_image(self):
        # print(str(self.all_img_paths[self.current_ind]))
        img = Image.open(str(self.all_img_paths[self.current_ind]))
        W, H = img.size
        img_ratio = W / H
        win_size = (self.lb_image.winfo_width(), self.lb_image.winfo_height())
        win_ratio = win_size[0] / win_size[1]
        # show_img = Image.fromarray(image)
        if win_ratio >= img_ratio:
            img_h = win_size[1]
            img_w = int(img_h * img_ratio)
        else:
            img_w = win_size[0]
            img_h = int(img_w / img_ratio)
        
        show_img = img.resize((img_w, img_h))
        imgtk = ImageTk.PhotoImage(show_img)
        self.lb_image.imgtk = imgtk
        self.lb_image.configure(image=imgtk)

        if self.all_labels[self.current_ind] is not None:
            self.list_all.itemconfig(self.current_ind, background=self.selected_colors[self.all_labels[self.current_ind]])
            cur_p = str(self.rel_paths[self.current_ind])
            if cur_p not in self.to_be_deleted_inds:
                self.to_be_deleted_inds.append(cur_p)
                self.rel_paths_set.remove(cur_p)
                # if self.fresh_rest_list:
                # rest_paths = eval(self.contents_rest_list.get())
                # self.contents_rest_list.set(r' '.join([p for p in rest_paths if p not in self.to_be_deleted_inds]))
        else:
            self.list_all.itemconfig(self.current_ind, background='white')
        self.list_all.see(self.current_ind)
        self.list_all.activate(self.current_ind)

        if self.all_labels[self.current_ind] is not None:
            self.lb_image['bg'] = self.selected_colors[self.all_labels[self.current_ind]]
            cur_cls = self.classnames[self.all_labels[self.current_ind]]
        else:
            self.lb_image['bg'] = 'white'
            cur_cls = '未标注'
        self.lf_image['text'] = "图像以及标注 {} @{} {:d} / {:d}  标注完成 {:.2f}%".format(
            str(self.rel_paths[self.current_ind]),
            cur_cls, 
            self.current_ind + 1,
            self.num_img,
            100 * len(self.to_be_deleted_inds) / self.num_img
        )




if __name__ == '__main__':
    master = tk.Tk()
    if os.name in ['nt']:
        master.iconbitmap('mv.ico')
    else:
        pass
    myapp = App(master)
    # print(font.families(myapp))
    myapp.master.title("PICKER v0.2 by Alpha@MOTOVIS")
    myapp.mainloop()
